
from major.models_manager import embedding_model, chat_model
from langchain_chroma import Chroma


'''
基于上一步构建好的知识点文档库，做一个知识询问rag
'''

# 加载文档库
vector_db = Chroma(persist_directory='./chroma_db',
                   embedding_function=embedding_model.get_model())

llm = chat_model.get_model()


# 检索与生成
def generate(query):
    # 检索一下与query相关的文档
    docs = vector_db.similarity_search(query, k=3)
    docs = [f"知识点{i + 1}:\n\n[[{doc.page_content}]]" for i, doc in enumerate(docs)]
    doc_context_str = "\n\n".join(docs)

    # 拼接prompt
    prompt = f'''
    请根据相关文档回答问题，如果文档中没有相关信息，请回答"没有相关信息"
    相关文档：
    {doc_context_str}

    用户问题：{query}

    '''
    res = llm.stream(prompt)
    for i in res:
        print(i.content, end='')


generate("什么是自回归生成")